# -*- coding: utf-8 -*-
from flask import Flask, render_template, request, jsonify
from auth import auth_login, require_login
from models import *
from PIL import Image

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = './'
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpeg', 'jpg','bmp','PNG', 'JPEG', 'JPG','BMP'}

def check_image(file):
    try:
        img = Image.open(file)
        file.seek(0, os.SEEK_END)
        file_size = file.tell()
        file.seek(0)
        width, height = img.size
        if width > 2048 or height > 2048:
            return False, "图像的长边不应超过2048像素"
        return True, None
    except Exception as e:
        return False, str(e)

def allowed_file(filename):
    """检查文件扩展名是否允许"""
    return '.' in filename and \
        filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']


@app.route('/login', methods=['POST', 'GET'])
def login():
    return auth_login()


@app.route('/', methods=['GET'])
@require_login
def index():
    return render_template('index.html')


@app.route('/api/upload', methods=['POST'])
@require_login
def upload_file():
    try:
        if 'file' not in request.files:
            return api_response(False, message="未找到文件部分", status_code=409)
        file = request.files['file']
        if file.filename == '':
            return api_response(False, message="没有选择文件", status_code=409)
        if file and allowed_file(file.filename):
            valid, error_msg = check_image(file)

            if not valid:
                return api_response(False, message=error_msg, status_code=400)
            filename = file.filename
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(filepath)
            vi_url = request_segmentation(filepath)
            os.remove(filepath)
            return api_response(True, data={"url": vi_url})
        else:
            return api_response(False, message="文件类型不支持", status_code=409)
    except Exception as e:
        return api_response(False, message=f"服务器内部错误: {str(e)}", status_code=500)


@app.route('/api/generation', methods=['POST'])
@require_login
def generation():
    try:
        data = request.json
        template_image_url = data.get('template_image_url')
        ref_prompt = data.get('ref_prompt')
        valid, results = generate_background(template_image_url, ref_prompt)
        if not valid:
            return api_response(False, message=results, status_code=400)
        return api_response(True, data=results)
    except Exception as e:
        return api_response(False, message=f"服务器内部错误: {str(e)}", status_code=500)

def api_response(status, message='', data=None, status_code=200):
    """
    构造API响应的辅助函数。

    :param status: 响应状态（布尔值或字符串：例如 True/False 或 'success', 'error'）
    :param message: 响应消息（可选）
    :param data: 响应数据（可选，默认为空字典）
    :param status_code: HTTP状态码（默认200）
    :return: JSON格式的响应和HTTP状态码
    """
    # 根据布尔值转换status为相应的字符串
    if isinstance(status, bool):
        status = 'success' if status else 'error'

    response = {
        'status': status,
        'message': message,
        'data': data if data is not None else {}
    }
    return jsonify(response), status_code


if __name__ == '__main__':
    app.run(host="0.0.0.0", port=9000)
